# -*- coding: utf-8 -*-

# 为了完成这个任务，我们可以编写一个Python类，使用ArcPy库来处理DEM（数字高程模型）数据。这个类将接收DEM文件路径、降雨量以及输出目录作为输入，并计算洼地范围、淹没深度、淹没面积和淹没体积。以下是一个可能的实现：

import arcpy
import os
from arcpy import env
from arcpy.sa import *


class FloodAnalysis:
    def __init__(self, dem_path, rainfall, output_dir):
        self.dem_path = dem_path
        self.rainfall = rainfall
        self.output_dir = output_dir
        self.flooded_areas = {}

    def _create_output_directory(self):
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)

    def _fill_sinks(self, dem):
        filled_dem = FillSinks(dem, "SORT")
        return filled_dem

    def _calculate_flow_direction(self, filled_dem):
        flow_dir = FlowDirection(filled_dem)
        return flow_dir

    def _calculate_flow_accumulation(self, flow_dir):
        flow_acc = FlowAccumulation(flow_dir)
        return flow_acc

    def _determine_flood_limits(self, flow_acc, rainfall):
        flood_limits = flow_acc <= (self.dem_path.split(".")[0] + "_flood_limits")
        arcpy.Conditional_analysis(flow_acc, rainfall, flood_limits, "LESS_EQUAL")
        return flood_limits

    def _calculate_flooded_areas(self, dem, flood_limits):
        flooded_areas = dem * flood_limits
        return flooded_areas

    def _calculate_flood_depth(self, dem, flooded_areas):
        flood_depth = flooded_areas - dem
        return flood_depth

    def _calculate_flood_volume(self, flood_depth):
        flood_volume = ZonalStatisticsAsTable(flood_depth, "VALUE", "SUM", cellSize=1)
        return flood_volume[0][0]

    def _calculate_flood_area(self, flooded_areas):
        flood_area = arcpy.GetRasterProperties_management(flooded_areas, "COUNT")
        return flood_area.getOutput(0)

    def _save_results(self, flooded_areas, flood_depth, flood_area, flood_volume):
        self.flooded_areas[self.rainfall] = {
            "flooded_areas": flooded_areas,
            "flood_depth": flood_depth,
            "flood_area": flood_area,
            "flood_volume": flood_volume
        }
        for key, raster in self.flooded_areas[self.rainfall].items():
            output_raster_path = os.path.join(self.output_dir, key + ".tif")
            arcpy.SaveRaster_management(raster, output_raster_path)

        output_shp_path = os.path.join(self.output_dir, "flooded_areas.shp")
        arcpy.FeatureToPolygon_management(flooded_areas, output_shp_path, cell_size=1)

    def run_analysis(self):
        self._create_output_directory()

        dem = Raster(self.dem_path)
        filled_dem = self._fill_sinks(dem)
        flow_dir = self._calculate_flow_direction(filled_dem)
        flow_acc = self._calculate_flow_accumulation(flow_dir)
        flood_limits = self._determine_flood_limits(flow_acc, self.rainfall)
        flooded_areas = self._calculate_flooded_areas(dem, flood_limits)
        flood_depth = self._calculate_flood_depth(dem, flooded_areas)
        flood_area = self._calculate_flood_area(flooded_areas)
        flood_volume = self._calculate_flood_volume(flood_depth)

        self._save_results(flooded_areas, flood_depth, flood_area, flood_volume)

    # 示例用法


fa = FloodAnalysis(r"path_to_dem.tif", 100, r"path_to_output_directory")
fa.run_analysis()
